from PIL import Image
from tensorbay import GAS
from tensorbay.dataset import Dataset
import numpy as np
import csv
import mindspore.dataset as ds

def getDataSet():

    BATCH_NUM = 4
    BATCH_SIZE = 1

    print("---------------------    start loading dataset    ---------------------")

    gas = GAS("Accesskey-ee75f1427dbca536bedc3bbec492f5ce")

    dataset = Dataset("RP2K", gas)

    train_segment = dataset["train"]

    train_segment = train_segment[:(BATCH_NUM * BATCH_SIZE)]

    hash_map = {}

    csv_file = csv.reader(open('lables_csv.csv', 'r'))

    for pair in csv_file:
        hash_map[pair[0]] = int(pair[1])

    datas = []
    lables = []

    for testcase in train_segment:
        with testcase.open() as fp:
            image = Image.open(fp)
            image = image.resize((32, 32))
            im_array = np.array(image)
            datas.append(im_array)
            lables.append(hash_map[testcase.label.classification.category])

    class DatasetGenerator:
        def __init__(self):
            self.image = np.array(datas)
            self.label = np.array(lables)

        def __getitem__(self, index):
            return self.image[index], self.label[index]

        def __len__(self):
            return len(self.image)

    dataset_generator = DatasetGenerator()
    dataset = ds.GeneratorDataset(dataset_generator, ["image", "label"], shuffle=False)
    print("--------------------- loading dataset successfully ---------------------")
    return dataset





